import numpy as np


class Particle:
    """
    粒子类，用于粒子群优化算法

    Attributes:
        _min_lims (array-like, optional): 粒子位置的最小限制，即待优化参数的下边界
        _max_lims (array-like, optional): 粒子位置的最大限制，即待优化参数的上边界
        _position (numpy.ndarray): 粒子的当前位置，即待优化参数的当前尝试值
        _velocity (numpy.ndarray): 粒子的当前速度，即迭代过程中用于改变待优化参数的值
        _best_position (numpy.ndarray): 粒子的历史最佳位置，即历史尝试中的最优参数
        _best_cost (float): 粒子的历史最佳代价，即历史尝试中的最小代价

    Methods:
        __init__(self, _dimension, _min_lims=None, _max_lims=None):
            类的构造方法，初始化粒子对象
        update_particle_velo_and_pos(self, _global_best_position, _inertial_weight, _cognitive_weight, _social_weight):
            在所有粒子完成一轮迭代后，根据各个粒子历史和本轮产生的最佳代价，更新粒子的速度和位置
        update_particle_best(self, _current_cost):
            在当前粒子使用当前位置完成代价计算后，更新粒子的历史最佳位置和代价，如果当前代价更大，则不更新
        get_position(self):
            从外部获取粒子的当前位置
        get_best(self):
            从外部获取粒子的历史最佳位置和代价
    """

    def __init__(self, _dimension, _min_lims=None, _max_lims=None):
        """
        类的构造方法，初始化粒子对象

        Parameters:
            _dimension (int): 粒子的维度（待优化参数的个数）
            _min_lims (array-like, optional): 粒子位置的最小限制
            _max_lims (array-like, optional): 粒子位置的最大限制
        """

        self._min_lims = _min_lims
        self._max_lims = _max_lims
        self._position = np.random.uniform(low=_min_lims,
                                           high=_max_lims,
                                           size=_dimension)
        self._velocity = np.random.uniform(low=_min_lims,
                                           high=_max_lims,
                                           size=_dimension)
        self._best_position = np.copy(self._position)
        self._best_cost = float('inf')

    def update_particle_velo_and_pos(self, _global_best_position,
                                     _inertial_weight, _cognitive_weight,
                                     _social_weight):
        """
        在所有粒子完成一轮迭代后，根据各个粒子历史和本轮产生的最佳代价，更新粒子的速度和位置

        Parameters:
            _global_best_position (numpy.ndarray): 全局最佳位置，由所有粒子循环一次后更新得到
            _inertial_weight (float): 本轮迭代惯性权重，由群类计算并给定
            _cognitive_weight (float): 本轮迭代认知权重，由群类计算并给定
            _social_weight (float): 本轮迭代社会权重，由群类计算并给定
        """

        inertial_term = _inertial_weight * self._velocity
        cognitive_term = _cognitive_weight * np.random.rand() * (
            self._best_position - self._position)
        social_term = _social_weight * np.random.rand() * (
            _global_best_position - self._position)

        self._velocity = inertial_term + cognitive_term + social_term
        self._position = self._position + self._velocity

        self._position = np.clip(self._position, self._min_lims,
                                 self._max_lims)

    def update_particle_best(self, _current_cost):
        """
        在当前粒子使用当前位置完成代价计算后，更新粒子的历史最佳位置和代价，如果当前代价更大，则不更新

        Parameters:
            _current_cost (float): 当前代价
        """

        if _current_cost < self._best_cost:
            self._best_cost = _current_cost
            self._best_position = np.copy(self._position)

    def get_position(self):
        """
        从外部获取粒子的当前位置

        Returns:
            numpy.ndarray: 粒子的当前位置
        """

        return self._position

    def get_best(self):
        """
        从外部获取粒子的历史最佳位置和代价

        Returns:
            tuple: 包含历史最佳位置和代价的元组
        """

        return self._best_position, self._best_cost


class Swarm:
    """
    （粒子）群类，用于粒子群优化算法

    Attributes:
        _current_particle (int): 在本次迭代中，当前处理的粒子的索引
        _swarm_size (int): 粒子群规模
        _iteration_size (int): 总迭代次数
        _particles (list): 包含多个粒子对象的列表，大小为粒子群的规模
        _global_best_position (numpy.ndarray): 全局最佳位置
        _global_best_cost (float): 全局最佳代价

    Methods:
        __init__(self, _dimension, _swarm_size, _iteration_size, _min_lims=None, _max_lims=None):
            类的构造方法，初始化（粒子）群对象
        _update_weights(self, _current_iteration):
            根据当前迭代轮次，更新粒子群权重参数，目前为线性权重参数改变
        update_swarm(self, _current_iteration):
            更新粒子群的状态，先根据每个粒子的最优来更新全局最优，后更新每一个粒子的下一轮迭代参数
        get_particle_position(self):
            从外部获取当前操作的粒子的位置
        return_particle_best(self, _current_cost):
            更新当前粒子的历史最佳位置和代价
        get_swarm_best(self):
            从外部获取粒子群的历史最佳位置和代价
    """

    def __init__(self,
                 _dimension,
                 _swarm_size,
                 _iteration_size,
                 _min_lims=None,
                 _max_lims=None):
        """
        类的构造方法，初始化（粒子）群对象

        Parameters:
            _dimension (int): 单个粒子的维度
            _swarm_size (int): 粒子群规模
            _iteration_size (int): 总迭代次数
            _min_lims (array-like, optional): 粒子位置的最小限制
            _max_lims (array-like, optional): 粒子位置的最大限制
        """

        self._current_particle = 0
        self._swarm_size = _swarm_size
        self._current_iteration = 0
        self._iteration_size = _iteration_size
        self._particles = [
            Particle(_dimension, _min_lims, _max_lims)
            for _ in range(self._swarm_size)
        ]
        self._global_best_position = None
        self._global_best_cost = float('inf')

    def _update_weights(self):
        """
        根据当前迭代轮次，更新粒子群权重参数，目前为线性权重参数改变

        Returns:
            tuple: 包含惯性权重、认知权重、社会权重的元组
        """

        inertial_weight = 0.9 - (0.5 / (self._iteration_size - 1)) * (
            self._current_iteration - 1)
        cognitive_weight = 2.5 - (2.0 / (self._iteration_size - 1)) * (
            self._current_iteration - 1)
        social_weight = 0.5 + (2.0 / (self._iteration_size - 1)) * (
            self._current_iteration - 1)
        return inertial_weight, cognitive_weight, social_weight

    def update_swarm(self):
        """
        更新粒子群的状态，先根据每个粒子的最优来更新全局最优，后更新每一个粒子的下一轮迭代参数
        """

        _current_best_cost = self._global_best_cost
        for particle in self._particles:
            b_pos, b_cost = particle.get_best()
            if b_cost < self._global_best_cost:
                self._global_best_cost = b_cost
                self._global_best_position = np.copy(b_pos)

        self._current_iteration = self._current_iteration + 1

        i_weight, c_weight, s_weight = self._update_weights()
        print("weight parameters for %d iteration: %.2f, %.2f, %.2f" %
              (self._current_iteration, i_weight, c_weight, s_weight))
        for particle in self._particles:
            particle.update_particle_velo_and_pos(self._global_best_position,
                                                  i_weight, c_weight, s_weight)

    def get_particle_position(self):
        """
        从外部获取当前粒子的位置

        Returns:
            numpy.ndarray: 当前粒子的位置
        """
        if self._current_particle >= (self._swarm_size - 1):
            self._current_particle = -1
        self._current_particle = self._current_particle + 1
        return self._current_particle, self._particles[
            self._current_particle].get_position()

    def return_particle_best(self, _current_particle, _current_cost):
        """
        更新当前粒子的历史最佳位置和代价

        Parameters:
            _current_particle (int): 当前计算完成的粒子编号
            _current_cost (float): 当前代价
        """

        self._particles[_current_particle].update_particle_best(_current_cost)

    def get_swarm_best(self):
        """
        从外部获取粒子群的历史最佳位置和代价

        Returns:
            tuple: 包含历史最佳位置和代价的元组
        """

        return self._global_best_position, self._global_best_cost


if __name__ == '__main__':
    '''
    def schwefel_function(_position):
        _length = len(_position)
        _sum_term = np.sum(-_position * np.sin(np.sqrt(np.abs(_position))))
        _result = 418.9829 * _length + _sum_term
        return _result
    '''

    def compute_cost(_position):
        return (10 * _position[0] * _position[1] +
                _position[1] * _position[2] + _position[2] * _position[3] +
                _position[3] / _position[4] + _position[4] / _position[5] +
                _position[5] / _position[0])

    dimension = 6
    swarm_size = 20 * dimension
    iteration_size = 100
    low_limit = 0.1
    high_limit = 5
    min_lims = np.array([low_limit] * dimension)
    max_lims = np.array([high_limit] * dimension)

    particle_swarm = Swarm(dimension, swarm_size, iteration_size, min_lims,
                           max_lims)

    for ite in range(0, iteration_size):
        for n in range(0, swarm_size):
            particle_number, particle_position = particle_swarm.get_particle_position(
            )
            current_cost = compute_cost(particle_position)
            particle_swarm.return_particle_best(particle_number, current_cost)

        particle_swarm.update_swarm()

        position, cost = particle_swarm.get_swarm_best()
        position = [round(x, 3) for x in position]
        print("best position:", position, "best cost: {: .7f}".format(cost))
